from langchain.llms.base import LLM
from typing import Optional, List, Dict, Any
import requests
from dotenv import load_dotenv
import os

# 加载环境变量（存储MindIE接口地址，避免代码硬编码）
load_dotenv()

class Qwen2MindIELLM(LLM):
    """适配MindIE部署的Qwen2模型的LangChain自定义LLM类"""
    # 1. 可配置参数（与MindIE接口参数对应，用户可调整）
    model_name: str = "DeepSeekR1"  # 模型标识（自定义）
    mindie_url: str = os.getenv("MINDIE_QLM_URL")  # 从.env读取接口地址
    max_tokens: int = 512  # 最大输出长度
    temperature: float = 0.7  # 随机性（0-1，越低越稳定）

    # 2. 核心方法：调用MindIE接口获取Qwen2回答
    def _call(
        self,
        prompt: str,  # LangChain传入的提示词（必须参数）
        stop: Optional[List[str]] = None,  # 停止词（MindIE不支持可忽略）
        run_manager: Optional[Any] = None,
    ) -> str:
        # 构造MindIE接口需要的请求参数（需与你的接口参数完全匹配）
        request_data = {
            "prompt": prompt,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature
            # 若你的接口有其他参数（如top_p），需在此补充
        }

        # 发送HTTP请求调用Qwen2
        try:
            response = requests.post(
                url=self.mindie_url,
                json=request_data,
                headers={"Content-Type": "application/json"},
                timeout=30  # Qwen2生成可能慢，设足够超时时间
            )
            response.raise_for_status()  # 若接口返回4xx/5xx，抛出异常
            result = response.json()  # 解析接口响应
        except Exception as e:
            return f"Qwen2调用失败：{str(e)}"  # 异常提示

        # 提取模型回答（需与你的接口响应字段匹配，示例是"output"）
        if result.get("code") != 200:  # 假设code=200为成功
            return f"模型生成失败：{result.get('msg', '未知错误')}"
        return result.get("output", "Qwen2未返回有效内容")

    # 3. 必须实现：返回模型标识参数（LangChain内部用）
    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {
            "model_name": self.model_name,
            "mindie_url": self.mindie_url,
            "max_tokens": self.max_tokens,
            "temperature": self.temperature
        }

    # 4. 必须实现：返回模型类型（自定义，标识为Qwen2）
    @property
    def _llm_type(self) -> str:
        return "qwen2-mindie"
